import torch
from PIL import Image
from torchvision import transforms as T

""" 模拟三通道二维卷积 """
img = Image.open("../../images/lena.png")
transformer = T.Compose([
    T.Resize((100, 100), antialias=True),
    T.ToTensor()
])
img = transformer(img)
print(img.shape)    # torch.Size([3, 100, 100]) CHW


# 定义3个3*3的卷积核
# 为什么是3个呢？每个通道都有自己独属的卷积核

kernel = torch.tensor([
    # 第一个 3*3
    [[-1, 0, 1],
    [-2, 0, 2],
    [-1, 0, 1]],
    # 第二个
    [[1, 0, 1],
    [2, -1, 0],
    [1, 0, -1]],
    # 第三个
    [[2, 0, 1],
    [1, 0, 2],
    [1, -1, 0]],
])
print(kernel.shape) # torch.Size([3, 3, 3])



# 模拟整体的卷积
C, H, W = img.shape
KC, KH, KW = kernel.shape
out_channel = 3
conv_img = torch.zeros(out_channel, H-KH+1, W-KW+1)
for c_out in range(out_channel):     # 遍历每个输出通道
    for i in range(H-KH+1):
        for j in range(W-KW+1):
            t = (img[:, i:i+3, j:j+3] * kernel[c_out]).sum()
            # 单个输出通道的结果，由所有输入通道与对应卷积核通道的加权和得到的
            conv_img[c_out, i, j] = t

print(conv_img.shape)

# 进行归一化
conv_img = (conv_img - conv_img.min()) / (conv_img.max() - conv_img.min())

T.ToPILImage()(conv_img).show()

"""
通过显示卷积后的图像，我们可以理解：卷积的作用就是从图像中提取特征（边缘、纹理等）

在上面的卷积中，我们虽然定义了conv_img是3通道，但实际我们给三通道赋的值是同一个，那么就显示为灰色
若想要彩色的，就需要有 3个 3*3*3的卷积核

也就是卷积核的维度 = (输出通道数, 输入通道数, 高度, 宽度)
"""